import os
import numpy as np
import pickle
import time
import torch
import torch.nn as nn
import torch.utils.data
from torchvision import transforms, datasets
from torch.utils.data import random_split, DataLoader
from collections import OrderedDict
import matplotlib.pyplot as plt  # 导入matplotlib
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

# 定义多层感知器
class TwoLayerNet:
    def __init__(self, input_size, hidden_size, output_size, weight_init_std=0.01):
        self.params = {}
        self.params['W1'] = weight_init_std * np.random.randn(input_size, hidden_size)
        self.params['b1'] = np.zeros(hidden_size)
        self.params['W2'] = weight_init_std * np.random.randn(hidden_size, output_size)
        self.params['b2'] = np.zeros(output_size)

        self.layers = OrderedDict()
        self.layers['Affine1'] = Affine(self.params['W1'], self.params['b1'])
        self.layers['Relu1'] = Relu()
        self.layers['Affine2'] = Affine(self.params['W2'], self.params['b2'])
        self.lastLayer = SoftmaxWithLoss()

    def predict(self, x):
        for layer in self.layers.values():
            x = layer.forward(x)
        return x

    def loss(self, x, t):
        y = self.predict(x)
        return self.lastLayer.forward(y, t)

    def accuracy(self, x, t):
        y = self.predict(x)
        y = np.argmax(y, axis=1)
        if t.ndim != 1: t = np.argmax(t, axis=1)
        accuracy = np.sum(y == t) / float(x.shape[0])
        return accuracy

    def numerical_gradient(self, x, t):
        loss_W = lambda W: self.loss(x, t)
        grads = {}
        grads['W1'] = numerical_gradient(loss_W, self.params['W1'])
        grads['b1'] = numerical_gradient(loss_W, self.params['b1'])
        grads['W2'] = numerical_gradient(loss_W, self.params['W2'])
        grads['b2'] = numerical_gradient(loss_W, self.params['b2'])
        return grads

    def gradient(self, x, t):
        self.loss(x, t)
        dout = 1
        dout = self.lastLayer.backward(dout)

        layers = list(self.layers.values())
        layers.reverse()
        for layer in layers:
            dout = layer.backward(dout)

        grads = {}
        grads['W1'], grads['b1'] = self.layers['Affine1'].dW, self.layers['Affine1'].db
        grads['W2'], grads['b2'] = self.layers['Affine2'].dW, self.layers['Affine2'].db
        return grads


# 定义损失和激活函数类
class Affine:
    def __init__(self, W, b):
        self.W = W
        self.b = b

    def forward(self, x):
        self.x = x
        return np.dot(x, self.W) + self.b

    def backward(self, dout):
        self.dW = np.dot(self.x.T, dout)
        self.db = np.sum(dout, axis=0)
        dx = np.dot(dout, self.W.T)
        return dx


class Relu:
    def forward(self, x):
        self.mask = (x <= 0)
        out = x.copy()
        out[self.mask] = 0
        return out

    def backward(self, dout):
        dout[self.mask] = 0
        dx = dout
        return dx


class SoftmaxWithLoss:
    def forward(self, x, t):
        self.t = t
        self.y = softmax(x)
        return cross_entropy_error(self.y, self.t)

    def backward(self, dout=1):
        batch_size = self.t.shape[0]
        if self.t.size == self.y.size:  # 如果t是one-hot向量
            dx = (self.y - self.t) / batch_size
        else:
            dx = self.y.copy()
            dx[np.arange(batch_size), self.t] -= 1
            dx = dx / batch_size
        return dx


def softmax(x):
    if x.ndim == 2:
        x = x - x.max(axis=1, keepdims=True)
        exp_x = np.exp(x)
        return exp_x / exp_x.sum(axis=1, keepdims=True)
    x = x - np.max(x)  # 为了防止溢出
    return np.exp(x) / np.sum(np.exp(x))


def cross_entropy_error(y, t):
    if y.ndim == 1:
        t = t.reshape(1, t.size)
        y = y.reshape(1, y.size)
    if t.size == y.size:  # 如果t是one-hot向量
        t = t.argmax(axis=1)
    batch_size = y.shape[0]
    return -np.sum(np.log(y[np.arange(batch_size), t] + 1e-7)) / batch_size


def numerical_gradient(f, x):
    h = 1e-4
    grad = np.zeros_like(x)
    it = np.nditer(x, flags=['multi_index'], op_flags=['readwrite'])

    while not it.finished:
        idx = it.multi_index
        tmp_val = x[idx]
        x[idx] = float(tmp_val) + h
        fxh1 = f(x)
        x[idx] = tmp_val - h
        fxh2 = f(x)
        grad[idx] = (fxh1 - fxh2) / (2 * h)
        x[idx] = tmp_val
        it.iternext()
    return grad


# 定义LeNet-5网络
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 6, 5),
            nn.Sigmoid(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 5),
            nn.Sigmoid(),
            nn.MaxPool2d(2, 2)
        )
        self.fc = nn.Sequential(
            nn.Linear(16 * 4 * 4, 120),
            nn.Sigmoid(),
            nn.Linear(120, 84),
            nn.Sigmoid(),
            nn.Linear(84, 10)
        )

    def forward(self, img):
        feature = self.conv(img)
        output = self.fc(feature.view(img.shape[0], -1))
        return output


# 加载 MNIST 数据集
transform = transforms.Compose([transforms.ToTensor()])
dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)

# 按照 75% 训练集和 25% 验证集划分数据集
train_size = int(0.75 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = random_split(dataset, [train_size, val_size])

train_iter = DataLoader(train_dataset, batch_size=100, shuffle=True)
val_iter = DataLoader(val_dataset, batch_size=100, shuffle=False)

# 初始化网络
device = 'cuda' if torch.cuda.is_available() else 'cpu'
lr, num_epochs = 0.001, 3
net = LeNet().to(device)
optimizer = torch.optim.Adam(net.parameters(), lr=lr)
loss = torch.nn.CrossEntropyLoss()

# 记录训练和验证过程中的损失和精度
train_losses = []
train_accuracies = []
val_losses = []
val_accuracies = []

# 训练过程
for i in range(num_epochs):
    train_l_sum, train_acc_sum, n, batch_count, start = 0.0, 0.0, 0, 0, time.time()

    # 训练阶段
    net.train()  # 设置为训练模式
    for step, (batch_x, batch_y) in enumerate(train_iter):
        batch_x = batch_x.to(device)
        batch_y = batch_y.to(device)

        y_hat = net(batch_x)
        loss_ = loss(y_hat, batch_y)

        optimizer.zero_grad()
        loss_.backward()
        optimizer.step()

        train_l_sum += loss_.item()
        train_acc_sum += (y_hat.argmax(dim=1) == batch_y).sum().item()
        batch_count += 1
        n += batch_y.shape[0]

    train_losses.append(train_l_sum / batch_count)
    train_accuracies.append(train_acc_sum / n)

    # 验证阶段
    val_l_sum, val_acc_sum, val_n = 0.0, 0.0, 0
    net.eval()  # 设置为评估模式
    with torch.no_grad():  # 不计算梯度
        for val_x, val_y in val_iter:
            val_x = val_x.to(device)
            val_y = val_y.to(device)

            val_y_hat = net(val_x)
            val_loss = loss(val_y_hat, val_y)

            val_l_sum += val_loss.item()
            val_acc_sum += (val_y_hat.argmax(dim=1) == val_y).sum().item()
            val_n += val_y.shape[0]

    val_losses.append(val_l_sum / len(val_iter))
    val_accuracies.append(val_acc_sum / val_n)

    print('epoch %d, loss %.4f, train acc %.3f, val acc %.3f, time %.1f sec' % (
        i + 1, train_l_sum / batch_count, train_acc_sum / n, val_acc_sum / val_n, time.time() - start))

# 绘制训练和验证的损失和准确率曲线
plt.figure(figsize=(12, 5))

# 绘制损失曲线
plt.subplot(1, 2, 1)
plt.plot(range(1, num_epochs + 1), train_losses, label='Train Loss')
plt.plot(range(1, num_epochs + 1), val_losses, label='Validation Loss')
plt.title('Loss Curve')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()

# 绘制准确率曲线
plt.subplot(1, 2, 2)
plt.plot(range(1, num_epochs + 1), train_accuracies, label='Train Accuracy')
plt.plot(range(1, num_epochs + 1), val_accuracies, label='Validation Accuracy')
plt.title('Accuracy Curve')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()

plt.tight_layout()
plt.show()

# 评估测试集
test_data = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
test_iter = DataLoader(test_data, batch_size=100, shuffle=False)

test_acc_sum, n = 0, 0
y_true = []
y_pred = []

sample_img = None  # 用于保存一张测试图片
sample_label = None  # 用于保存对应的标签

for test_x, test_y in test_iter:
    test_x = test_x.to(device)
    test_y = test_y.to(device)
    y_hat = net(test_x)

    test_acc_sum += (y_hat.argmax(dim=1) == test_y).sum().item()
    n += test_y.shape[0]

    # 记录真实标签和预测结果
    y_true.extend(test_y.cpu().numpy())
    y_pred.extend(y_hat.argmax(dim=1).cpu().numpy())

    # 获取第一张测试样本进行显示
    sample_img = test_x[0].cpu().detach()  # 从GPU转移到CPU
    sample_label = test_y[0].cpu().detach()  # 从GPU转移到CPU
    predicted_label = y_hat[0].argmax(dim=0).cpu().detach()  # 预测标签

    # 只需显示一张图片，跳出循环
    break

print('Test Accuracy: %.3f' % (test_acc_sum / n))

# 绘制混淆矩阵
confusion_mtx = confusion_matrix(y_true, y_pred)
disp = ConfusionMatrixDisplay(confusion_matrix=confusion_mtx, display_labels=np.arange(10))
disp.plot(cmap=plt.cm.Blues)
plt.title('Confusion Matrix')
plt.show()

# 显示测试样本图片
plt.imshow(sample_img.squeeze(), cmap='gray')
plt.title(f'True: {sample_label.item()}, Predicted: {predicted_label.item()}')
plt.axis('off')
plt.show()
